from a2a.server.request_handlers.request_handler import RequestHandler
from agent import BackgroundAgent
from a2a.server.context import ServerCallContext
from a2a.types import (
    Message,
    MessageSendParams,
    Task,
    TaskIdParams,
    TaskPushNotificationConfig,
    TaskQueryParams,
    TaskStatus,
    TaskState,
    TaskStatusUpdateEvent,
    Role,
    TextPart,
)
from collections.abc import AsyncGenerator
from a2a.server.events.event_queue import Event
from typing import Optional, Any, Dict
import uuid


class BackgroundAgentHandler(RequestHandler):
    def __init__(self, agent:BackgroundAgent):
        self.agent = agent

    async def on_get_task(self, params: TaskQueryParams, context: ServerCallContext | None = None) -> Optional[Task]:
        return None

    async def on_cancel_task(self, params: TaskIdParams, context: ServerCallContext | None = None) -> Optional[Task]:
        return None

    async def on_message_send(
            self, 
            params: MessageSendParams, 
            context: ServerCallContext | None = None
        ) -> Task:
        # 创建新任务消息：hello world
        task_message = await self.agent.invoke()

        # 返回一个任务对象
        return Task(
            id=params.message.taskId or str(uuid.uuid4()),
            contextId=params.message.contextId or str(uuid.uuid4()),
            status=TaskStatus(
                state=TaskState.completed,
                message=Message(
                    messageId=str(uuid.uuid4()),
                    role=Role.agent,
                    kind="message",
                    parts=[TextPart(text=task_message)],
                    contextId=params.message.contextId,
                    taskId=params.message.taskId
                )
            )
        )
        
    async def on_message_send_stream(
            self, 
            params: MessageSendParams, 
            context: ServerCallContext | None = None
        ) -> AsyncGenerator[Event,None]:

        # 创建新任务消息：hello world
        task_message = await self.agent.invoke()
        print (f"Task message: {task_message}")
        
        # 使用简单的方式创建消息，避免参数问题
        yield Message(
            messageId=str(uuid.uuid4()),
            role=Role.agent,
            kind="message",
            parts=[TextPart(text=task_message)],
            contextId=params.message.contextId,
            taskId=params.message.taskId
        )
        # 任务结束更新状态
        yield TaskStatusUpdateEvent(
            contextId=params.message.contextId or "",
            taskId=params.message.taskId or "",
            final=True,
            kind="status-update",
            status=TaskStatus(
                state=TaskState.completed
            )
        )

    async def on_set_task_push_notification_config(
        self, 
        params: TaskPushNotificationConfig,
        context: ServerCallContext | None = None
    ) -> TaskPushNotificationConfig:
        pass

    async def on_get_task_push_notification_config(
        self, 
        params: TaskIdParams,
        context: ServerCallContext | None = None
    ) -> TaskPushNotificationConfig:
        pass

    async def on_resubscribe_to_task(
        self, 
        params: TaskIdParams,
        context: ServerCallContext | None = None
    ) -> AsyncGenerator[Event]:
        pass